package empty.dsa.generator


import chisel3._
import chisel3.util._
import empty.dsa.util.alu_Info.sha1_penumMap

class sha1_alu(elename: String, width: Int, pe_num: Int, NO_PE: Int) extends Module {

  override val desiredName = elename
  val io = IO(new Bundle {
    val cfg = Input(UInt(1.W)) //通过cfg判断为循环的起始位置
    val inputs = Input(Vec(7, UInt(width.W))) //a,b,c,d,e,cycle,w
    val outputs = Output(Vec(6, UInt(width.W)))
  })
  val inputsWire = Wire(Vec(6, UInt(width.W)))
  val cycle = Wire(UInt(7.W))
  if(pe_num == 80) {
    if (NO_PE == 1) {
      inputsWire(0) := "h67452301".U
      inputsWire(1) := "hEFCDAB89".U
      inputsWire(2) := "h98badcfe".U
      inputsWire(3) := "h10325476".U
      inputsWire(4) := "hC3D2E1F0".U
      cycle := 0.U
      inputsWire(5) := RegNext(io.inputs(6))
    } else {
      for (i <- 0 until 5) {
        inputsWire(i) := RegNext(io.inputs(i))
      }
        inputsWire(5) := RegNext(io.inputs(6))
        cycle := RegNext(io.inputs(5)(6,0))
    }
  }
  else if(NO_PE == 1) {
    when(io.cfg === 1.U) {
      inputsWire(0) := "h67452301".U
      inputsWire(1) := "hEFCDAB89".U
      inputsWire(2) := "h98badcfe".U
      inputsWire(3) := "h10325476".U
      inputsWire(4) := "hC3D2E1F0".U
      cycle := 0.U
      inputsWire(5) := RegNext(io.inputs(6))
    }.otherwise {
      for (i <- 0 until 5) {
        inputsWire(i) := RegNext(io.inputs(i))
      }
        inputsWire(5) := RegNext(io.inputs(6))
        cycle := RegNext(io.inputs(5)(6,0))

    }
  } else {
    for (i <- 0 until 5) {
      inputsWire(i) := RegNext(io.inputs(i))

    }
    inputsWire(5) := RegNext(io.inputs(6))
    cycle := RegNext(io.inputs(5)(6,0))
  }

  val kt = Wire(UInt(32.W))
  val k_ori = Seq("h5A827999".U, "h6ED9EBA1".U, "h8F1BBCDC".U, "hCA62C1D6".U)
  val k = VecInit(Seq.fill(sha1_penumMap(pe_num)(0))(0.U(32.W)))


  if(pe_num == 80) {
    if(NO_PE <= 20) {
      k := Seq(k_ori(0))
    } else if(20 < NO_PE && NO_PE <= 40) {
      k := Seq(k_ori(1))
    }  else if(40 < NO_PE && NO_PE <= 60) {
      k := Seq(k_ori(2))
    } else {
      k := Seq(k_ori(3))
    }
  } else if(pe_num == 40) {
    if(NO_PE < 21) {
      k := Seq(k_ori(0), k_ori(2))
    } else {
      k := Seq(k_ori(1), k_ori(3))
    }
  } else {
    k := k_ori
  }

  //根据PE_num选择kt
  if(pe_num == 80) {
    kt := k(0)
  } else if(pe_num == 40) { // 二选一
    kt := Mux(cycle < 40.U, k(0), k(1))
  } else { //四选一
    kt := MuxCase(0.U, Array(
      (cycle < 20.U) -> k(0),
      (20.U <= cycle && cycle < 40.U) -> k(1),
      (40.U <= cycle && cycle < 60.U) -> k(2),
      (60.U <= cycle && cycle < 80.U) -> k(3)))
  }

  val sha1func_res = Wire(UInt(width.W))
  val sha1func_res0 = Wire(UInt(width.W))
  val sha1func_res1 = Wire(UInt(width.W))
  val sha1func_res2 = Wire(UInt(width.W))


  if(pe_num == 80){
    if(NO_PE <= 20) {
      sha1func_res := sha1func_res0
    } else if(20 < NO_PE && NO_PE <= 40) {
      sha1func_res := sha1func_res1
    }  else if(40 < NO_PE && NO_PE <= 60) {
      sha1func_res := sha1func_res2
    } else {
      sha1func_res := sha1func_res1
    }
  } else if(pe_num == 40) {
    if(NO_PE < 21) {
      sha1func_res := Mux(cycle > 39.U, sha1func_res2, sha1func_res0)
    } else {
      sha1func_res := sha1func_res1
    }
  } else {
    sha1func_res := MuxCase(0.U, Array(
      (inputsWire(5) <= 19.U) -> sha1func_res0,
      (19.U < cycle & cycle <= 39.U) -> sha1func_res1,
      (39.U < cycle & cycle <= 59.U) -> sha1func_res2,
      (59.U < cycle & cycle <= 79.U) -> sha1func_res1))
  }


  //(b&c) | (~b&d)
  sha1func_res0 := (inputsWire(1) & inputsWire(2)) | ((~inputsWire(1)) & inputsWire(3))
  //b^c^d
  sha1func_res1 := inputsWire(1) ^ inputsWire(2) ^ inputsWire(3)
  //(b&c) | (b&d) | (c&d)
  sha1func_res2 := (inputsWire(1) & inputsWire(2)) | (inputsWire(1) & inputsWire(3)) | (inputsWire(2) & inputsWire(3))

  val shift_a = Wire(UInt(width.W))
  val shift_b = Wire(UInt(width.W))
  val tmp_A = Wire(UInt(width.W))
  shift_a := Cat(inputsWire(0)(26, 0), inputsWire(0)(31, 27))
  shift_b := Cat(inputsWire(1)(1, 0), inputsWire(1)(31, 2))
  tmp_A := shift_a + sha1func_res + inputsWire(4) + kt + inputsWire(5)

  if(NO_PE == 80) {
    io.outputs(0) := tmp_A + "h67452301".U
    io.outputs(1) := inputsWire(0) + "hEFCDAB89".U
    io.outputs(2) := shift_b + "h98badcfe".U
    io.outputs(3) := inputsWire(2) + "h10325476".U
    io.outputs(4) := inputsWire(3) + "hC3D2E1F0".U
    io.outputs(5) := 0.U
  } else if(NO_PE == pe_num) {
    when(cycle === 79.U) {
      io.outputs(0) := tmp_A + "h67452301".U
      io.outputs(1) := inputsWire(0) + "hEFCDAB89".U
      io.outputs(2) := shift_b + "h98badcfe".U
      io.outputs(3) := inputsWire(2) + "h10325476".U
      io.outputs(4) := inputsWire(3) + "hC3D2E1F0".U
      io.outputs(5) := 0.U
    } .otherwise{
      io.outputs(0) := tmp_A
      io.outputs(1) := inputsWire(0)
      io.outputs(2) := shift_b
      io.outputs(3) := inputsWire(2)
      io.outputs(4) := inputsWire(3)
      io.outputs(5) := cycle + 1.U
    }
  } else {
    io.outputs(0) := tmp_A
    io.outputs(1) := inputsWire(0)
    io.outputs(2) := shift_b
    io.outputs(3) := inputsWire(2)
    io.outputs(4) := inputsWire(3)
    io.outputs(5) := cycle + 1.U
  }
}

//object sha1Gen extends App {
//  chisel3.Driver.execute(args, () => new sha1_alu("alu_sha1", 32, 16, 1))
//}